import torch
import torchvision as torchvision
from torch import nn
from net import MyLeNet5  # 导入搭建的网络
from torch.optim import lr_scheduler  # 学习率 优化器
from torchvision import datasets, transforms  # 数据集加载
from torch.utils.data import DataLoader, Dataset
import os

# os.getcwd()获取当前路径

# 数据转换为tensor格式
data_transform = transforms.Compose([
    transforms.ToTensor()
])

# 加载内置的训练数据集：minst
train_dataset = datasets.MNIST(root="./data", train=True, transform=data_transform, download=True)
train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
# 加载内置的测试数据集：minst
test_dataset = datasets.MNIST(root="./data", train=False, transform=data_transform, download=True)
test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, shuffle=True)

# 如果有显卡，可以转到GPU
device = "cuda" if torch.cuda.is_available() else 'cpu'

# 调用之前搭建好的网络模型，将模型数据转到GPU
model = MyLeNet5().to(device)

# 定义一个损失函数：交叉熵损失函数
loss_fn = nn.CrossEntropyLoss()

# 定义一个优化器 SGD随机梯度下降法==》模型参数进行反向传播 lr学习率  momentum控制下降速度
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

# 学习率更新：每隔10轮，变为原来的0.1
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)


# 定义一个训练函数
def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0
    for _, (X, y) in enumerate(dataloader):  # (X,y)X表示图片，用numpy——》tensor表示
        # 前向传播
        X, y = X.to(device), y.to(device)  # 将数据丢入GPU/CPU中
        output = model(X)  # 经过网络结构之后的结果 (16,10)
        cur_loss = loss_fn(output, y)  # 计算出”差“
        _, pred = torch.max(output, axis=1)  # 输出每行的最大值，预测的最大可能classes

        cur_acc = torch.sum(y == pred) / output.shape[0]  # 得到每一batch的精确度

        # 更新optimizer、cur_loss、optimizer
        optimizer.zero_grad()  # 梯度置零，也就是把loss关于weight的导数变成0.
        cur_loss.backward()  # 反向传播求梯度
        optimizer.step()  # 对应weights = [weights[k] + alpha * d_weights[k] for k in range(n)]

        loss += cur_loss.item()  # 把一个元素tensor当值显示出来,类似于value,算出总loss
        current += cur_acc.item()  # 计算总的精度
        n = n + 1   # 计算出总的循环次数，即训练样本数
    print("train_loss" + str(loss / n))
    print("train_acc" + str(current / n))


# 定义一个验证函数，参数不更新
def val(dataloader, model, loss_fn):
    model.eval()    # 将模型进行验证，进行验证模式
    loss, current, n = 0.0, 0.0, 0
    with torch.no_grad():   # 固定进入，强制之后的内容不进行计算图构建。不进行梯度计算
        for _, (X, y) in enumerate(dataloader):
            # 前向传播
            X, y = X.to(device), y.to(device)
            output = model(X)
            cur_loss = loss_fn(output, y)
            _, pred = torch.max(output, axis=1)
            cur_acc = torch.sum(y == pred) / output.shape[0]
            loss += cur_loss.item()
            current += cur_acc.item()
            n = n + 1
        print("val_loss" + str(loss / n))
        print("val_acc" + str(current / n))

        return current / n


# 开始训练
epoch = 50  # 训练轮次
min_acc = 0  # 最小精度

for t in range(epoch):
    print(f"epoch{t + 1}\n---------------------------")
    train(train_dataloader, model, loss_fn, optimizer)  # 进行训练
    a = val(test_dataloader, model, loss_fn)    # 进行验证 ，无优化器optimizer

    # 只保存最好的模型权重
    if a > min_acc:
        folder = "save_model"
        if not os.path.exists(folder):
            os.mkdir(folder)
        min_acc = a
        print("save best model")
        torch.save(model.state_dict(), 'save_model/best_model.pth')
    print("\n")

print("Done!")

# torch.load(url) 加载打开权重文件
